package arts.spbu.testswing;

import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * User: pev67021
 * Date: 08.11.12
 * Time: 10:12
 */

// Префиксное дерево

public class Trie {

    private static Map<String, ArrayList<WordAttrs>> lemmaMap = new HashMap<>();

    static class TrieNode {
        Map<Character, TrieNode> children = new HashMap<>();
        boolean leaf = false;
        ArrayList<String> lems = new ArrayList<>();
    }

    private TrieNode root = new TrieNode();

    public void put(String s, List<String> lemmas) {
        TrieNode v = root;
        int pos = Integer.parseInt(s.substring(0, s.indexOf("_")));
        s = s.substring(s.indexOf("_") + 1);
        for (char ch : s.toLowerCase().toCharArray()) {
            if (!v.children.containsKey(ch))
                v.children.put(ch, new TrieNode());
            v = v.children.get(ch);
        }
        v.leaf = true;

        for (String lemma : lemmas) {
            String lem = lemma;
            String attrs = "";
            if (lemma.contains("=")) {
                lem = lemma.substring(0, lemma.indexOf("="));
                attrs = lemma.substring(lemma.indexOf("="));
            }

            if (lem.endsWith("?"))
                lem = lem.substring(0, lem.indexOf("?"));

            if (lemmaMap.containsKey(lem)) {
                //System.out.println("LEMMA: " + lem + " attrs: " + attrs);
                lemmaMap.get(lem).add(new WordAttrs(s, attrs, pos));
            } else {
                ArrayList<WordAttrs> temp = new ArrayList<>();
                temp.add(new WordAttrs(s, attrs, pos));
                //System.out.println("LEMMA: " + lem + " attrs: " + attrs);
                lemmaMap.put(lem, temp);
            }
            v.lems.add(lem);
        }
    }

    public boolean find(String s) {
        TrieNode v = root;
        for (char ch : s.toLowerCase().toCharArray()) {
            if (!v.children.containsKey(ch))
                return false;
            else
                v = v.children.get(ch);
        }
        return v.leaf;
    }

    public String[] getLemmas(String s) {
        TrieNode v = root;
        for (char ch : s.toLowerCase().toCharArray()) {
            if (!v.children.containsKey(ch))
                return null;
            else
                v = v.children.get(ch);
        }
        String[] lemmas = new String[v.lems.size()];
        for (int i = 0; i < v.lems.size(); i++)
            lemmas[i] = v.lems.get(i);
        return lemmas;
    }

    public ArrayList<String> getForms(String lemma) {
        if (!lemmaMap.containsKey(lemma))
            return null;
        else {
            ArrayList<String> wforms = new ArrayList<>();
            for (WordAttrs WdAttrs : lemmaMap.get(lemma))
                wforms.add(WdAttrs.getWord());
            return wforms;
        }
    }

    public ArrayList<String> getAttrsByPos(String lemma, int pos) {
        if (!lemmaMap.containsKey(lemma))
            return null;
        else {
            ArrayList<String> wforms = new ArrayList<>();
            for (WordAttrs WdAttrs : lemmaMap.get(lemma)) {
                //System.out.println("Looking though wordform "+ WdAttrs.getWord() + " at " + WdAttrs.getStartPos());
                if (WdAttrs.getStartPos() == pos)
                    wforms.add(WdAttrs.getAttrs());
            }
            return wforms;
        }
    }

    public boolean findLemma(String lemma) {
        return lemmaMap.containsKey(lemma);
    }

    public ArrayList<String> getWFormAttrsByPos(String wordform, int pos) {
        ArrayList<String> attrs = new ArrayList<>();
        ArrayList<String> temp_attrs;
        for (String lemma : this.getLemmas(wordform)) {
            //System.out.println("Current lemma " + lemma);
            temp_attrs = this.getAttrsByPos(lemma, pos);
            //System.out.println("And what about the attrs?");
            if (temp_attrs == null) {
                //System.out.println("No attributes for " + wordform + " at " + pos);
                return null;
            } else
                //System.out.println("NOT NULL attrs");
                for (String attr : temp_attrs) {
                    //System.out.println("temp attrs: " + attr);
                    attrs.addAll(Arrays.asList(attr.split("\\|")));
                }
            //System.out.println("That's all about the attrs");
        }
        return attrs;
    }

}
